package com.boat.framework.common.servlet;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintWriter;

import javax.servlet.ServletOutputStream;
import javax.servlet.ServletResponse;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;

import org.apache.commons.io.output.TeeOutputStream;

/**
 * 响应包装器。用于servlet打印日志
 * 
 * @author 周江
 *
 */
public class ResponseWrapper extends HttpServletResponseWrapper {

	private final ByteArrayOutputStream bos = new ByteArrayOutputStream();
	private PrintWriter writer = new PrintWriter(bos);

	public ResponseWrapper(HttpServletResponse response) {
		super(response);
	}

	@Override
	public ServletResponse getResponse() {
		return this;
	}

	@Override
	public ServletOutputStream getOutputStream() throws IOException {
		return new ServletOutputStream() {
			@Override
			public boolean isReady() {
				return false;
			}

			@Override
			public void setWriteListener(WriteListener listener) {

			}

			private TeeOutputStream tee = new TeeOutputStream(ResponseWrapper.super.getOutputStream(), bos);

			@Override
			public void write(int b) throws IOException {
				tee.write(b);
			}
		};
	}

	@Override
	public PrintWriter getWriter() throws IOException {
		return new TeePrintWriter(super.getWriter(), writer);
	}

	public byte[] toByteArray() {
		return bos.toByteArray();
	}

	/**
	 * PrintWriter是一种写入字符的一种操作类，可以写入字符，TeePrintWriter继承了他，主要功能是把原始的字符流copy到branch里面。
	 * 
	 * @author 周江
	 *
	 */
	public class TeePrintWriter extends PrintWriter {

		PrintWriter branch;

		public TeePrintWriter(PrintWriter main, PrintWriter branch) {
			super(main, true);
			this.branch = branch;
		}

		public void write(char buf[], int off, int len) {
			super.write(buf, off, len);
			super.flush();
			branch.write(buf, off, len);
			branch.flush();
		}

		public void write(String s, int off, int len) {
			super.write(s, off, len);
			super.flush();
			branch.write(s, off, len);
			branch.flush();
		}

		public void write(int c) {
			super.write(c);
			super.flush();
			branch.write(c);
			branch.flush();
		}

		public void flush() {
			super.flush();
			branch.flush();
		}
	}
}